package com.easy.secret.codec;

import java.io.UnsupportedEncodingException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Security;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.pqc.math.linearalgebra.ByteUtils;
import org.springframework.util.Assert;

/**
 * @ProjectName: esay-secret-server
 * @Package: com.easy.secret.codec
 * @ClassName: SM4Codec
 * @Description: 国密4算法加解密
 * @Author: wangjc
 * @CreateDate: 2023/11/17
 * @UpdateUser: wangjc
 * @UpdateDate: 2023/11/17
 * @UpdateRemark: 更新内容
 * @Version: 1.0
 */
public class SM4Codec {
	public static String CHAR_SET;
	public static final String RANDOM_CHARS = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!@#$%^&*";
	public static final String ALGORITHM_NAME = "SM4";
	private String key;
	private String iv;
	private String type = "ECB";
	private String padding = "SM4/ECB/PKCS7Padding";
	private String toStringType = "base64";

	public SM4Codec(String key, String iv, String type, String padding, String toStringType) {
		this.key = key;
		Assert.hasText(key, "sm4加密得秘钥不能为空");
		if (StringUtils.isNotBlank(iv)) {
			this.iv = iv;
		}

		if (StringUtils.isNotBlank(type)) {
			this.type = type;
		}

		if (StringUtils.isNotBlank(padding)) {
			this.padding = padding;
		}

		if (StringUtils.isNotBlank(toStringType)) {
			this.toStringType = toStringType;
		}
	}

	public static String generateKey() {
		return RandomStringUtils.random(16, RANDOM_CHARS);
	}

	public static String getIvParameter() {
		return RandomStringUtils.random(16, RANDOM_CHARS);
	}

	public String encrypt(String data) throws Exception {
		return this.encrypt(data, (String) null);
	}

	public String encrypt(String data, String charset) throws Exception {
		if (StringUtils.isNotBlank(charset)) {
			CHAR_SET = charset;
		}

		if ("hex".equalsIgnoreCase(this.toStringType)) {
			return ByteUtils.toHexString(this.encrypt(data.getBytes(CHAR_SET)));
		} else if ("base64".equalsIgnoreCase(this.toStringType)) {
			return Base64Codec.instance().encode(this.encrypt(data.getBytes(CHAR_SET)));
		} else {
			throw new RuntimeException("encrypt error");
		}
	}

	byte[] encrypt(byte[] data) throws Exception {
		Cipher cipher;
		if (this.type.equalsIgnoreCase("ECB")) {
			cipher = generateECBCipher(this.padding, 1, this.key.getBytes(CHAR_SET));
			return cipher.doFinal(data);
		} else if (this.type.equalsIgnoreCase("CBC")) {
			cipher = generateCBCCipher(this.padding, 1, this.key.getBytes(CHAR_SET), this.iv.getBytes(CHAR_SET));
			return cipher.doFinal(data);
		} else {
			throw new RuntimeException("encrypt error");
		}
	}

	public String decrypt(String data) throws Exception {
		return this.decrypt(data, (String) null);
	}

	public String decrypt(String data, String charset) throws Exception {
		if (StringUtils.isNotBlank(charset)) {
			CHAR_SET = charset;
		}

		byte[] decrypt = null;
		if ("hex".equalsIgnoreCase(this.toStringType)) {
			decrypt = this.decrypt(ByteUtils.fromHexString(data));
		} else {
			if (!"base64".equalsIgnoreCase(this.toStringType)) {
				throw new RuntimeException("decrypt error");
			}
			decrypt = this.decrypt(Base64Codec.instance().decodeByte(data, CHAR_SET));
		}
		return new String(decrypt);
	}

	byte[] decrypt(byte[] data) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException,
			NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException,
			InvalidAlgorithmParameterException, UnsupportedEncodingException {
		Cipher cipher;
		if (this.type.equalsIgnoreCase("ECB")) {
			cipher = generateECBCipher(this.padding, 2, this.key.getBytes(CHAR_SET));
			return cipher.doFinal(data);
		} else if (this.type.equalsIgnoreCase("CBC")) {
			cipher = generateCBCCipher(this.padding, 2, this.key.getBytes(CHAR_SET), this.iv.getBytes(CHAR_SET));
			return cipher.doFinal(data);
		} else {
			throw new RuntimeException("decrypt error");
		}
	}

	private static Cipher generateECBCipher(String algorithmName, int mode, byte[] key)
			throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException {
		Cipher cipher = Cipher.getInstance(algorithmName, "BC");
		Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
		cipher.init(mode, sm4Key);
		return cipher;
	}

	private static Cipher generateCBCCipher(String algorithmName, int mode, byte[] key, byte[] iv)
			throws InvalidKeyException, InvalidAlgorithmParameterException, NoSuchAlgorithmException,
			NoSuchProviderException, NoSuchPaddingException {
		Cipher cipher = Cipher.getInstance(algorithmName, "BC");
		Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
		IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
		cipher.init(mode, sm4Key, ivParameterSpec);
		return cipher;
	}

	static {
		Security.addProvider(new BouncyCastleProvider());
		CHAR_SET = "UTF-8";
	}
}
